package socket

import (
	"sync"
)

var Sub *SubStationAlarm

func init() {
	Sub = newSubStationAlarm()
}

type SubStationAlarm struct {
	rw  sync.RWMutex
	sub map[string]map[int64]interface{}
}

func newSubStationAlarm() *SubStationAlarm {
	return &SubStationAlarm{
		sub: make(map[string]map[int64]interface{}),
	}
}

// SubStationMsg 订阅站点消息.
func (s *SubStationAlarm) SubStationMsg(userId int64, stationId string) {
	s.rw.Lock()
	defer s.rw.Unlock()
	if _, ok := s.sub[stationId]; !ok {
		s.sub[stationId] = make(map[int64]interface{})
	}
	s.sub[stationId][userId] = 1
}

// UnSubStationMsg 取消订阅.
func (s *SubStationAlarm) UnSubStationMsg(userId int64, stationId string) {
	s.rw.Lock()
	defer s.rw.Unlock()
	if _, ok := s.sub[stationId]; ok {
		delete(s.sub[stationId], userId)
	}
}

// GetSubUserIds 获取订阅的用户ID.
func (s *SubStationAlarm) GetSubUserIds(stationId string) []int64 {
	s.rw.RLock()
	defer s.rw.RUnlock()
	if _, ok := s.sub[stationId]; ok {
		resp := make([]int64, 0, len(s.sub[stationId]))
		for id, _ := range s.sub[stationId] {
			resp = append(resp, id)
		}
		return resp
	}
	return []int64{}
}

// UnSubAll 离开所有.
func (s *SubStationAlarm) UnSubAll(userId int64) {
	s.rw.Lock()
	defer s.rw.Unlock()
	for k := range s.sub {
		delete(s.sub[k], userId)
	}
}
